import pickle

import process_efw

def store_model(picklefile = None, alphabet = None, gram_size = 4):
    if alphabet is None:
        alphabet = ['#','a','b','c','d','e','f','g','h','i','j','k','l',\
                    'm','n','o','p','q','r','s','t','u','v','w','x','y',\
                    'z','-', "'", '.']

    words = process_efw.process_file("my_efw")
    grammar = process_efw.WordCorpusGram(words, alphabet)
    model = calc_model(grammar, gram_size)

    if picklefile is None:
        picklefile = open("charmodel", 'w')

    pickle.dump((gram_size, model), picklefile)


def load_model(picklefile = None):
    if picklefile is None:
        picklefile = open("charmodel", "r")

    model = pickle.load(picklefile)
    picklefile.close()
    return model


def setup_gramstrings(alphabet, gramsize):

    grams = [""]
    for s in xrange(gramsize):
        newgrams = []

        for a in alphabet:
            for g in grams:
                newgrams.append(g + a)

        grams = newgrams

    return grams




def calc_model(wcgrammar, gramsize):
    gramstrings = setup_gramstrings(wcgrammar.alphabet, gramsize-1)
    grams = {}
    for g in gramstrings:
#        d = dict((a, 0.0) for a in wcgrammar.alphabet)
#        d[None] = 0
        grams[g] = [0.0 for a in wcgrammar.alphabet] + [0]

    for word in wcgrammar.words:
        pword = wcgrammar.alphabet[0]*(gramsize-1) + word
        for w in xrange(gramsize-1, len(pword)):
            context = pword[w-gramsize+1:w]
            grams[context][wcgrammar.alphabet.index(pword[w])] += 1
#            grams[context][pword[w]] += 1
            grams[context][-1] += 1

    for g in grams:
        gram = grams[g]
#        total = gram[None]
        total = gram[-1]
        del gram[-1]
        if total > 0:
            for a in xrange(len(gram)):
                gram[a] = gram[a] / total

    return grams
